import netron
import onnx
import torch
from torch import nn


class ownModule(nn.Module):
    def __init__(self):
        super(ownModule, self).__init__()
        self.Conv1_1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=1, padding=0, stride=1)

        self.Conv3_3 = nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, padding=1, stride=1)
        self.Conv5_5 = nn.Conv2d(in_channels=32, out_channels=16, kernel_size=5, padding=2, stride=1)

        self.Conv3_3_ = nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, padding=1, stride=1)
        self.Conv5_5_ = nn.Conv2d(in_channels=32, out_channels=16, kernel_size=5, padding=2, stride=1)

        self.Conv_last = nn.Conv2d(in_channels=16, out_channels=3, kernel_size=3, padding=1, stride=1)

    def forward(self, Fi, Fi_1):
        # 1*1卷积
        Fi = self.Conv1_1(Fi)
        Fi_1 = self.Conv1_1(Fi_1)

        fi = torch.concat([Fi, Fi_1], dim=0)
        fi_1 = torch.concat([Fi_1, Fi], dim=0)

        fi = self.Conv3_3(fi)
        fi_1 = self.Conv5_5(fi_1)

        temp = torch.concat([fi, fi_1], dim=0)
        fi_1 = torch.concat([fi_1, fi], dim=0)
        fi = temp

        fi = self.Conv3_3_(fi)
        fi_1 = self.Conv5_5_(fi_1)

        output = fi * fi_1
        output = output + Fi + Fi_1
        output = self.Conv_last(output)
        return output


Fi = torch.rand((3, 128, 128))
Fi_1 = torch.rand((3, 128, 128))
model = ownModule()
outputs = model(Fi, Fi_1)

# pip install netron onnx
# 模型可视化，将模型导出为 ONNX 格式
torch.onnx.export(model, (Fi, Fi_1), './onnx_model.onnx')  # 导出后 netron.start(path) 打开
# 增加维度信息
onnx.save(onnx.shape_inference.infer_shapes(onnx.load("./onnx_model.onnx")), "./onnx_model.onnx")
# 打开导出的 ONNX 模型文件
netron.start('./onnx_model.onnx')
